import json
import os
def convert_file(input_path, output_path):
    """
    将原始的 instruction/input/output 格式转换为 Qwen 所需的 messages 格式
    """
    with open(input_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    converted = []
    for item in data:
        instruction = item.get("instruction", "").strip()
        input_content = item.get("input", "").strip()
        output = item.get("output", "").strip()
        user_content = instruction
        if input_content:
            user_content += "\n" + input_content
        messages = [
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": output}
        ]
        converted.append({"messages": messages})
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(converted, f, ensure_ascii=False, indent=4)
    print(f"✅ 已保存: {output_path} (共 {len(converted)} 条)")

def main(input_dir, output_dir):
    # 定义目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    files_to_convert = ["train.json", "test.json"]
    for filename in files_to_convert:
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)

        if not os.path.exists(input_path):
            print(f"⚠️ 警告: 文件不存在 {input_path}")
            continue
        print(f"🔄 正在处理 {input_path} ...")
        convert_file(input_path, output_path)
    print(f"\n🎉 所有文件已转换完成，输出到: {output_dir}/")
if __name__ == "__main__":
    input_dir = "./lora/NL2SQL_LORA"
    output_dir = f"{input_dir}_PRO"
    main(input_dir, output_dir)